/* Copyright by Kenneth Lee, 2024. All Rights Reserved */
#include <stdio.h>
#include <assert.h>
#include <string.h>

static void sched(void);
#define yield() sched()

int job1(void) {
	int i, sum;
	printf("job0\n");
	for(i=0,sum=0; i<100; i++) {
		yield();
		sum+=i;
	}
	return sum;
}

int job2(void) {
	int i, sum;
	printf("job1\n");
	for(i=0,sum=0; i<100; i++) {
		yield();
		sum+=i*2;
	}
	return sum;
}

int job3(void) {
	int i, sum;
	printf("job2\n");
	for(i=0,sum=0; i<100; i++) {
		yield();
		sum+=i*3;
	}
	return sum;
}

// 任务控制块(TCB）列表，为了简单我们这里固定支持10个任务
#define MAX_TCBS 10
#define STACK_SZ 0x10000
struct tcb {
	// job被挂起的时候的寄存器状态
	long    rsp,
		rax,
		rbx,
		rcx,
		rdx,
		rsi,
		rdi,
		rbp,
		r8,
		r9,
		r10,
		r11,
		r12,
		r13,
		r14,
		r15,
		rflags;

	long stack[STACK_SZ];    // job堆栈
	int valid;               // tcb是否正在使用，0表示未使用
	int to_exit;             // job准备退出，0表示不是准备退出
}jobs[MAX_TCBS];
struct tcb sys_job;              // 系统上下文，这个其实不需要堆栈，不过懒得拆开了，就浪费掉多余的堆栈吧。
int current_tcb = -1;

void job_exit(void) {
	// 退出前需要借用当前的堆栈调用sched()，所以不能马上去掉valid
	printf("job %d exit\n", current_tcb);
	jobs[current_tcb].to_exit = 1;
	sched();
}

void sched_create_job(int (*job)(void)) {
	int i;
	// 找一个空闲的tcb放新的任务
	for (i = 0; i < MAX_TCBS; i++) {
		if (jobs[i].valid == 0) {
			// 初始化
			memset(&jobs[i], 0, sizeof(jobs[i]));
			jobs[i].valid = 1;
			jobs[i].stack[STACK_SZ-1] = (long)job_exit;     // job退出地址压栈
			jobs[i].stack[STACK_SZ-2] = (long)job;          // job入口压栈
			jobs[i].rsp = (long)&jobs[i].stack[STACK_SZ-2]; // 设置栈顶
			jobs[i].rdi= (long)&jobs[i];                    // TCB指针（切换的第一个参数）
			return;
		}
	}
	assert(0); /* 我们假定不会创建超过数量的job */
}

extern void switch_to(long prev, long next);

// 这个算法的实现，就是所谓的调度算法
static long find_next(void) {
	int next;
	for (int i = 0; i < MAX_TCBS; i++) {
		next = (current_tcb + i + 1) % MAX_TCBS;
		if (jobs[next].valid) {
			if (jobs[next].to_exit)
				jobs[next].valid = 0;  // 这是准备退出的job，不调度，直接清除
			else
				return next;  // 这是下一个可以调度的进程
		}
	}

	return -1;
}

/*
static void dump_job(int i) {
	if (i >= 0) {
		printf("job %d (tcp: %p):\n", i, &jobs[i]);
		printf("\trax=%016lx rbx=%016lx rcx=%016lx rdx=%016lx\n",
			jobs[i].rax, jobs[i].rbx, jobs[i].rcx, jobs[i].rdx);
		printf("\trsi=%016lx rdi=%016lx rbp=%016lx rsp=%016lx\n",
			jobs[i].rsi, jobs[i].rdi, jobs[i].rbp, jobs[i].rsp);
		printf("\tr8=%016lx r9=%016lx r10=%016lx r11=%016lx\n",
			jobs[i].r8, jobs[i].r9, jobs[i].r10, jobs[i].r11);
		printf("\tr12=%016lx r13=%016lx r14=%016lx r15=%016lx\n",
			jobs[i].r12, jobs[i].r13, jobs[i].r14, jobs[i].r15);
		printf("\trflags=%016lx\n", jobs[i].rflags);
	}
}
*/

static void dump_switch_info(int from, int to) {
	printf("切换：%d -> %d\n", from, to);
	/*
	dump_job(from);
	dump_job(to);
	*/
}

// 进入调度
static void sched(void) {
	int next = find_next();
	if (next == -1) {
		printf("没有任务可以调度了，退出\n");
		switch_to(0, (long)&sys_job);
	} else if (next != current_tcb) {
		int prev = current_tcb;
		current_tcb = next;
		dump_switch_info(prev, next);
		switch_to((prev == -1) ? (long)&sys_job : (long)&jobs[prev],
			  (long)&jobs[next]);
	}
}

int main(void) {
	sched_create_job(job1);
	sched_create_job(job2);
	sched_create_job(job3);
	sched();

	return 0;
}
